from typing import Any, Optional
import json

from .config import Config
from .memory import Memory
from .utils.enum import ReportSource, ReportType, Tone
from .llm_provider import GenericLLMProvider
from .prompts import get_prompt_family
from .vector_store import VectorStoreWrapper

# Research skills
from .skills.researcher import ResearchConductor
from .skills.writer import ReportGenerator
from .skills.context_manager import ContextManager
from .skills.browser import BrowserManager
from .skills.curator import SourceCurator
from .skills.deep_research import DeepResearchSkill

from .actions import (
    add_references,
    extract_headers,
    extract_sections,
    table_of_contents,
    get_search_results,
    get_retrievers,
    choose_agent
)

"""
GPTResearcher 的单一代理实现
"""

"""
根据我看到的代码，GPTResearcher 是一个自主研究代理系统，能够对任何主题进行研究并生成带有引用的综合报告。
    
    以下是它的主要功能和组件：:-)

GPTResearcher 类是整个系统的核心，它包含以下关键组件：

1. 研究技能：
   - ResearchConductor：执行研究过程
   - ReportGenerator：生成报告
   - ContextManager：管理研究上下文
   - BrowserManager：处理网页浏览和抓取
   - SourceCurator：整理和验证来源
   - DeepResearchSkill：进行深度研究

2. 初始化参数：
   - query：研究查询
   - report_type：报告类型（研究报告或深度研究）
   - report_format：输出格式（默认为markdown）
   - report_source：报告来源（Web等）
   - tone：报告语调
   - 其他配置如来源URL、文档URL等

3. 主要方法：
   - conduct_research：执行研究过程
   - write_report：根据研究内容生成报告
   - write_report_conclusion：生成报告结论
   - write_introduction：生成介绍
   - quick_search：快速搜索信息
   - 其他工具方法：处理报告引用、提取标题、生成目录等

4. 特色功能：
   - 支持深度研究模式，可以更详细地探索主题
   - 可以处理多种数据源（网页、文档等）
   - 使用向量存储进行相似内容检索
   - 支持不同的报告格式和语调
   - 有完整的成本追踪和日志系统

系统使用LLM提供商（通过GenericLLMProvider）来处理自然语言任务，并使用各种技能模块来完成研究的不同阶段。:-)
"""

class GPTResearcher:
    def __init__(
        self,
        query: str,  # 研究查询/问题
        report_type: str = ReportType.ResearchReport.value,  # 报告类型：普通研究报告或深度研究
        report_format: str = "markdown",  # 报告输出格式，默认为markdown
        report_source: str = ReportSource.Web.value,  # 报告来源：Web、文档等
        tone: Tone = Tone.Objective,  # 报告语调：客观、批判性等
        source_urls: list[str] | None = None,  # 指定的源URL列表
        document_urls: list[str] | None = None,  # 文档URL列表
        complement_source_urls: bool = False,  # 是否补充源URL
        query_domains: list[str] | None = None,  # 限制查询的域名列表
        documents=None,  # 直接提供的文档
        vector_store=None,  # 向量存储，用于相似内容检索
        vector_store_filter=None,  # 向量存储过滤器
        config_path=None,  # 配置文件路径
        websocket=None,  # WebSocket连接，用于实时通信
        agent=None,  # 使用的代理类型
        role=None,  # 代理角色
        parent_query: str = "",  # 父查询（用于分层研究）
        subtopics: list | None = None,  # 子主题列表
        visited_urls: set | None = None,  # 已访问的URL集合
        verbose: bool = True,  # 是否输出详细日志
        context=None,  # 研究上下文，存储收集到的信息
        headers: dict | None = None,  # HTTP请求头
        max_subtopics: int = 5,  # 最大子主题数量
        log_handler=None,  # 日志处理器
        prompt_family: str | None = None,  # 提示词系列
    ):
        self.query = query
        self.report_type = report_type
        self.cfg = Config(config_path)  # 初始化配置
        self.cfg.set_verbose(verbose)
        self.llm = GenericLLMProvider(self.cfg)  # 初始化语言模型提供商
        self.report_source = report_source if report_source else getattr(self.cfg, 'report_source', None)
        self.report_format = report_format
        self.max_subtopics = max_subtopics
        self.tone = tone if isinstance(tone, Tone) else Tone.Objective
        self.source_urls = source_urls
        self.document_urls = document_urls
        self.complement_source_urls = complement_source_urls
        self.query_domains = query_domains or []
        self.research_sources = []  # 已抓取的源列表，包括标题、内容和图片
        self.research_images = []  # 已选择的研究图片列表
        self.documents = documents
        self.vector_store = VectorStoreWrapper(vector_store) if vector_store else None  # 初始化向量存储包装器
        self.vector_store_filter = vector_store_filter
        self.websocket = websocket
        self.agent = agent
        self.role = role
        self.parent_query = parent_query
        self.subtopics = subtopics or []
        self.visited_urls = visited_urls or set()  # 追踪已访问的URL
        self.verbose = verbose
        self.context = context or []  # 研究上下文，存储收集到的信息
        self.headers = headers or {}
        self.research_costs = 0.0  # 研究成本跟踪
        self.retrievers = get_retrievers(self.headers, self.cfg)  # 获取检索器
        self.memory = Memory(  # 初始化内存组件，用于嵌入和存储
            self.cfg.embedding_provider, self.cfg.embedding_model, **self.cfg.embedding_kwargs
        )
        self.log_handler = log_handler
        self.prompt_family = get_prompt_family(prompt_family or self.cfg.prompt_family, self.cfg)  # 获取提示词系列

        # 初始化各个组件
        self.research_conductor: ResearchConductor = ResearchConductor(self)  # 研究执行器
        self.report_generator: ReportGenerator = ReportGenerator(self)  # 报告生成器
        self.context_manager: ContextManager = ContextManager(self)  # 上下文管理器
        self.scraper_manager: BrowserManager = BrowserManager(self)  # 网页抓取管理器
        self.source_curator: SourceCurator = SourceCurator(self)  # 源管理器
        self.deep_researcher: Optional[DeepResearchSkill] = None  # 深度研究技能
        if report_type == ReportType.DeepResearch.value:
            self.deep_researcher = DeepResearchSkill(self)  # 如果是深度研究，初始化深度研究组件

    async def _log_event(self, event_type: str, **kwargs):
        """Helper method to handle logging events"""
        if self.log_handler:
            try:
                if event_type == "tool":
                    await self.log_handler.on_tool_start(kwargs.get('tool_name', ''), **kwargs)
                elif event_type == "action":
                    await self.log_handler.on_agent_action(kwargs.get('action', ''), **kwargs)
                elif event_type == "research":
                    await self.log_handler.on_research_step(kwargs.get('step', ''), kwargs.get('details', {}))

                # Add direct logging as backup
                import logging
                research_logger = logging.getLogger('research')
                research_logger.info(f"{event_type}: {json.dumps(kwargs, default=str)}")

            except Exception as e:
                import logging
                logging.getLogger('research').error(f"Error in _log_event: {e}", exc_info=True)

    async def conduct_research(self, on_progress=None):
        """
        执行研究过程的主要方法。
        根据配置的report_type，决定使用普通研究还是深度研究。
        会选择合适的代理和角色，然后调用研究执行器来进行实际研究。
        
        参数:
            on_progress: 可选的回调函数，用于报告研究进度
            
        返回:
            研究上下文（收集到的研究信息）
        """
        # 记录研究开始事件
        await self._log_event("research", step="start", details={
            "query": self.query,
            "report_type": self.report_type,
            "agent": self.agent,
            "role": self.role
        })

        # 如果是深度研究类型，使用深度研究技能处理
        if self.report_type == ReportType.DeepResearch.value and self.deep_researcher:
            return await self._handle_deep_research(on_progress)

        # 如果没有指定代理和角色，自动选择合适的代理
        if not (self.agent and self.role):
            await self._log_event("action", action="choose_agent")
            self.agent, self.role = await choose_agent(
                query=self.query,
                cfg=self.cfg,
                parent_query=self.parent_query,
                cost_callback=self.add_costs,
                headers=self.headers,
                prompt_family=self.prompt_family,
            )
            await self._log_event("action", action="agent_selected", details={
                "agent": self.agent,
                "role": self.role
            })

        # 记录研究进行中事件
        await self._log_event("research", step="conducting_research", details={
            "agent": self.agent,
            "role": self.role
        })
        # 使用研究执行器进行实际研究，并获取研究上下文
        self.context = await self.research_conductor.conduct_research()

        # 记录研究完成事件
        await self._log_event("research", step="research_completed", details={
            "context_length": len(self.context)
        })
        return self.context

    async def _handle_deep_research(self, on_progress=None):
        """
        处理深度研究执行和日志记录的方法。
        
        深度研究会更详细地探索主题，可能涉及多个子主题和层次，
        使用DeepResearchSkill组件来执行。
        
        参数:
            on_progress: 可选的回调函数，用于报告研究进度
            
        返回:
            深度研究的上下文（收集到的研究信息）
        """
        # 记录深度研究配置
        await self._log_event("research", step="deep_research_initialize", details={
            "type": "deep_research",
            "breadth": self.deep_researcher.breadth,
            "depth": self.deep_researcher.depth,
            "concurrency": self.deep_researcher.concurrency_limit
        })

        # 记录深度研究开始
        await self._log_event("research", step="deep_research_start", details={
            "query": self.query,
            "breadth": self.deep_researcher.breadth,
            "depth": self.deep_researcher.depth,
            "concurrency": self.deep_researcher.concurrency_limit
        })

        # 运行深度研究并获取上下文
        self.context = await self.deep_researcher.run(on_progress=on_progress)

        # 获取总研究成本
        total_costs = self.get_costs()

        # 记录深度研究完成及成本
        await self._log_event("research", step="deep_research_complete", details={
            "context_length": len(self.context),
            "visited_urls": len(self.visited_urls),
            "total_costs": total_costs
        })

        # 记录最终成本更新
        await self._log_event("research", step="cost_update", details={
            "cost": total_costs,
            "total_cost": total_costs,
            "research_type": "deep_research"
        })

        # 返回研究上下文
        return self.context

    async def write_report(self, existing_headers: list = [], relevant_written_contents: list = [], ext_context=None, custom_prompt="") -> str:
        """
        生成研究报告的主要方法。
        
        参数:
            existing_headers: 已存在的报告标题列表，用于保持报告结构一致性
            relevant_written_contents: 已写好的相关内容，可以被引用或整合
            ext_context: 外部研究上下文，如果不提供则使用内部上下文
            custom_prompt: 自定义提示词，用于指导报告生成
            
        返回:
            生成的研究报告正文（通常为Markdown格式）
        """
        # 记录开始写报告的事件
        await self._log_event("research", step="writing_report", details={
            "existing_headers": existing_headers,
            "context_source": "external" if ext_context else "internal"
        })

        # 调用报告生成器生成报告
        report = await self.report_generator.write_report(
            existing_headers=existing_headers,
            relevant_written_contents=relevant_written_contents,
            ext_context=ext_context or self.context,
            custom_prompt=custom_prompt
        )

        # 记录报告完成的事件
        await self._log_event("research", step="report_completed", details={
            "report_length": len(report)
        })
        return report

    async def write_report_conclusion(self, report_body: str) -> str:
        """
        生成研究报告结论的方法。
        
        参数:
            report_body: 报告主体内容，用于总结和得出结论
            
        返回:
            报告结论部分（通常为Markdown格式）
        """
        # 记录开始写结论的事件
        await self._log_event("research", step="writing_conclusion")
        # 调用报告生成器生成结论
        conclusion = await self.report_generator.write_report_conclusion(report_body)
        # 记录结论完成的事件
        await self._log_event("research", step="conclusion_completed")
        return conclusion

    async def write_introduction(self):
        """
        生成研究报告介绍的方法。
        
        返回:
            报告介绍部分（通常为Markdown格式）
        """
        # 记录开始写介绍的事件
        await self._log_event("research", step="writing_introduction")
        # 调用报告生成器生成介绍
        intro = await self.report_generator.write_introduction()
        # 记录介绍完成的事件
        await self._log_event("research", step="introduction_completed")
        return intro

    async def quick_search(self, query: str, query_domains: list[str] = None) -> list[Any]:
        """
        执行快速搜索，获取查询的搜索结果。
        
        参数:
            query: 搜索查询
            query_domains: 限制搜索的域名列表
            
        返回:
            搜索结果列表
        """
        return await get_search_results(query, self.retrievers[0], query_domains=query_domains)

    async def get_subtopics(self):
        """
        获取研究主题的子主题列表。
        由报告生成器根据主查询分析确定合适的子主题。
        
        返回:
            子主题列表
        """
        return await self.report_generator.get_subtopics()

    async def get_draft_section_titles(self, current_subtopic: str):
        """
        获取当前子主题的草稿章节标题。
        
        参数:
            current_subtopic: 当前子主题
            
        返回:
            章节标题列表
        """
        return await self.report_generator.get_draft_section_titles(current_subtopic)

    async def get_similar_written_contents_by_draft_section_titles(
        self,
        current_subtopic: str,
        draft_section_titles: list[str],
        written_contents: list[dict],
        max_results: int = 10
    ) -> list[str]:
        """
        根据草稿章节标题获取相似的已写内容。
        用于减少内容重复，提高报告一致性。
        
        参数:
            current_subtopic: 当前子主题
            draft_section_titles: 草稿章节标题列表
            written_contents: 已写内容列表
            max_results: 最大结果数量
            
        返回:
            相似内容列表
        """
        return await self.context_manager.get_similar_written_contents_by_draft_section_titles(
            current_subtopic,
            draft_section_titles,
            written_contents,
            max_results
        )

    # Utility methods
    def get_research_images(self, top_k=10) -> list[dict[str, Any]]:
        """获取研究过程中收集的图片，最多返回top_k个"""
        return self.research_images[:top_k]

    def add_research_images(self, images: list[dict[str, Any]]) -> None:
        """添加研究图片到图片列表"""
        self.research_images.extend(images)

    def get_research_sources(self) -> list[dict[str, Any]]:
        """获取所有研究来源"""
        return self.research_sources

    def add_research_sources(self, sources: list[dict[str, Any]]) -> None:
        """添加研究来源到来源列表"""
        self.research_sources.extend(sources)

    def add_references(self, report_markdown: str, visited_urls: set) -> str:
        """向报告添加引用/参考文献"""
        return add_references(report_markdown, visited_urls)

    def extract_headers(self, markdown_text: str) -> list[dict]:
        """从Markdown文本中提取标题"""
        return extract_headers(markdown_text)

    def extract_sections(self, markdown_text: str) -> list[dict]:
        """从Markdown文本中提取章节"""
        return extract_sections(markdown_text)

    def table_of_contents(self, markdown_text: str) -> str:
        """根据Markdown文本生成目录"""
        return table_of_contents(markdown_text)

    def get_source_urls(self) -> list:
        """获取所有已访问的URL列表"""
        return list(self.visited_urls)

    def get_research_context(self) -> list:
        """获取研究上下文"""
        return self.context

    def get_costs(self) -> float:
        """获取累计研究成本"""
        return self.research_costs

    def set_verbose(self, verbose: bool):
        """设置是否输出详细日志"""
        self.verbose = verbose

    def add_costs(self, cost: float) -> None:
        """
        添加研究成本
        
        参数:
            cost: 成本金额
            
        异常:
            ValueError: 如果成本不是数字类型
        """
        if not isinstance(cost, (float, int)):
            raise ValueError("Cost must be an integer or float")
        self.research_costs += cost
        if self.log_handler:
            self._log_event("research", step="cost_update", details={
                "cost": cost,
                "total_cost": self.research_costs
            })
